import numpy as np
from mindspore import nn


class MySets:
    def __init__(self, depth=4, idxes=None, num=10000) -> None:
        self.depth = depth
        x = np.linspace(0, 2 * np.pi, num, dtype=np.float32)
        self.y = np.reshape(
            np.sin(x) + np.random.rand(num).astype(np.float32) / 100, (-1, 1))
        self.idxes = idxes

    def __len__(self):
        return len(self.idxes) - self.depth

    def __getitem__(self, idx):
        idx1 = self.idxes[idx]

        x = self.y[idx1:idx1+self.depth]
        y = self.y[idx1+self.depth]

        return x, y


def split_set(test_size=0.3, num=10000, depth=4):
    idxes = list(range(num - depth))

    test_num = int(test_size * (num - depth))

    test_idxes = np.random.choice(idxes, test_num, replace=False)

    train_idxes = list(set(idxes) - set(test_idxes))

    return MySets(idxes=train_idxes, num=num), MySets(idxes=test_idxes, num=num)

class MAPE(nn.Metric):
    def __init__(self):
        super().__init__()
        self.clear()

    def clear(self):
        self._sample_num = 0
        self._error_num = 0.0

    @nn.rearrange_inputs
    def update(self, *inputs):
        y_pre = inputs[0].asnumpy()
        y = inputs[1].asnumpy()
        err = np.abs(y_pre - y / y)

        self._error_num += err.sum()
        self._sample_num += y.shape[0]

    def eval(self):
        return self._error_num / self._sample_num
